{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Action Angle Networks",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "Copyright 2022 Google LLC.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ],
      "metadata": {
        "id": "7qMi5_E-R9lu"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title License\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ],
      "metadata": {
        "cellView": "form",
        "id": "xzh-n15URpBL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Action-Angle Networks\n",
        "\n"
      ],
      "metadata": {
        "id": "G_2R-K2HirHJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "example_directory = 'action_angle_networks'\n",
        "editor_relpaths = ('configs/action_angle_flow.py', 'shm_simulation.py', 'models.py', 'train.py')\n",
        "\n",
        "repo, branch = 'https://github.com/google-research/google-research', 'master'\n",
        "# (If you run this code in Jupyter[lab], then you're already in the\n",
        "#  example directory and nothing needs to be done.)\n",
        "\n",
        "#@title Fetch source code\n",
        "if 'google.colab' in str(get_ipython()):\n",
        "  import os\n",
        "  os.chdir('/content')\n",
        "  # Download source repo from Github.\n",
        "  if not os.path.isdir('s'):\n",
        "    !git clone --depth=1 -b $branch $repo sourcerepo\n",
        "  # Copy example files & change directory.\n",
        "  example_root_path = f'/content/{example_directory}'\n",
        "  if not os.path.isdir(example_root_path):\n",
        "    os.makedirs(example_root_path)\n",
        "    !cp -r sourcerepo/$example_directory/* \"$example_root_path\"\n",
        "  os.chdir(example_root_path)\n",
        "  from google.colab import files\n",
        "  for relpath in editor_relpaths:\n",
        "    files.view(f'{example_root_path}/{relpath}')"
      ],
      "metadata": {
        "id": "LHqPslxDPwk5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!pwd"
      ],
      "metadata": {
        "id": "12FyNb8BRPCC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install -r requirements.txt"
      ],
      "metadata": {
        "id": "GfpVq4K0ilrx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Base Imports\n",
        "from typing import *\n",
        "import functools\n",
        "import sys\n",
        "import tempfile\n",
        "\n",
        "import collections\n",
        "import chex\n",
        "import numpy as np\n",
        "import jax\n",
        "import jax.numpy as jnp\n",
        "import flax.linen as nn\n",
        "from flax.core import frozen_dict\n",
        "from flax.training import train_state\n",
        "import optax\n",
        "import distrax\n",
        "import tensorflow as tf\n",
        "import ml_collections\n",
        "import matplotlib.pyplot as plt\n",
        "import matplotlib.animation as animation\n",
        "import matplotlib\n",
        "import sklearn.preprocessing\n",
        "matplotlib.rc('animation', html='jshtml')"
      ],
      "metadata": {
        "id": "8yR2UOi6oew1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Source Imports\n",
        "%load_ext autoreload\n",
        "%autoreload 2\n",
        "import chm_simulation\n",
        "import models\n",
        "import shm_simulation\n",
        "import train\n",
        "from configs import action_angle_flow"
      ],
      "metadata": {
        "id": "bPBQtMXLS7z2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%load_ext tensorboard"
      ],
      "metadata": {
        "id": "Y3_KeMenruNv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## The Hamiltonian\n",
        "\n",
        "The Hamiltonian describes the 'energy' of a system in terms of\n",
        "the positions $q$ and momenta $p$ of the particles.\n",
        "\n",
        "For example, in the case of a simple harmonic oscillator, we have:\n",
        "\n",
        "$$\n",
        "H(q, p) = \\frac{p^2}{2m} + \\frac{m\\omega^2q^2}{2}\n",
        "$$\n",
        "\n",
        "The time evolution of any system with Hamiltonian $H(q, p)$ is given by\n",
        "Hamilton's equations:\n",
        "\n",
        "$$\n",
        "\\begin{align}\n",
        "\\frac{dq}{dt} &= \\frac{\\partial H}{\\partial p} \\\\\n",
        "\\frac{dp}{dt} &= -\\frac{\\partial H}{\\partial q}\n",
        "\\end{align}\n",
        "$$\n",
        "\n",
        "and hence, the Hamiltonian is conserved across time:\n",
        "\n",
        "$$\n",
        "\\frac{dH}{dt} = \\frac{\\partial H}{\\partial q} \\frac{dq}{dt} + \\frac{\\partial H}{\\partial p} \\frac{dp}{dt} = 0 \n",
        "$$\n",
        "\n",
        "## Action-Angle Coordinates\n",
        "\n",
        "For integrable systems, there exists a mapping from canonical coordinates ($q$, $p$) to action-angle coordinates ($\\theta$, $I$).\n",
        "\n",
        "In these coordinates, the Hamiltonian is only a function of the actions $I$,\n",
        "not the angles $\\theta$:\n",
        "\n",
        "$$\n",
        "H(q, p) \\to H(\\theta, I) = K(I)\n",
        "$$\n",
        "\n",
        "and hence, by Hamilton's equations:\n",
        "\n",
        "$$\n",
        "\\frac{d\\theta}{dt} = \\frac{dH}{dI}.\n",
        "$$\n",
        "\n",
        "which is a constant across time as $I$ does not change with time:\n",
        "\n",
        "$$\n",
        "\\frac{dI}{dt} = -\\frac{dH}{d\\theta} = 0.\n",
        "$$\n",
        "\n",
        "Thus, the angles $\\theta$ evolve linearly, while the actions $I$ remain constant through time. This makes them very simple to model, requiring no integration to find the state of the system at any instant of time."
      ],
      "metadata": {
        "id": "oKNteSzeXaUk"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "To enable our exploration, we focus on the simple harmonic oscillator system, introduced above."
      ],
      "metadata": {
        "id": "Lasb33S4i0qQ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## The Simple Harmonic Oscillator\n",
        "\n",
        "The Hamiltonian for the simple harmonic oscillator is:\n",
        "$$\n",
        "H(q, p) = \\frac{p^2}{2m} + \\frac{m\\omega^2q^2}{2}\n",
        "$$\n",
        "\n",
        "Applying Hamilton's equations,\n",
        "$$\n",
        "\\begin{align}\n",
        "\\frac{dq}{dt} &= \\frac{\\partial H}{\\partial p} = \\frac{p}{m} \\\\\n",
        "\\frac{dp}{dt} &= -\\frac{\\partial H}{\\partial q} = -m\\omega^2q\n",
        "\\end{align}\n",
        "$$\n",
        "\n",
        "which can be solved as:\n",
        "$$\n",
        "\\begin{align}\n",
        "q &= A\\cos(\\omega t + \\phi) \\\\\n",
        "p &= −m\\omega A\\sin(\\omega t + \\phi)\n",
        "\\end{align}\n",
        "$$\n",
        "where $A$ and $\\phi$ are integration constants.\n",
        "\n",
        "This allows us to generate $(q, p)$ coordinates at any instant, given the values of $A, \\phi, m$ and $\\omega$."
      ],
      "metadata": {
        "id": "7ycFj5kbYTZV"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Parameters of system.\n",
        "simulation_parameters = {\n",
        "    'phi': jnp.asarray([0, jnp.pi/2]),\n",
        "    'A': jnp.asarray([1, 2]),\n",
        "    'm': jnp.asarray([1, 1]),\n",
        "    'w': jnp.asarray([0.05, 0.2]),\n",
        "}"
      ],
      "metadata": {
        "id": "xHYqqGEyPXyq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Generate coordinates at multiple instants of time!\n",
        "times = jnp.arange(100) * 1\n",
        "generate_canonical_coordinates_fn = jax.vmap(\n",
        "    jax.vmap(shm_simulation.generate_canonical_coordinates, in_axes=(0, None)),\n",
        "    in_axes=(None, 0), out_axes=1)\n",
        "positions, momentums = generate_canonical_coordinates_fn(times, simulation_parameters)"
      ],
      "metadata": {
        "id": "-Yb2dJSwXrBE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Compute Hamiltonians at the different instants.\n",
        "hamiltonians = jax.vmap(shm_simulation.compute_hamiltonian, in_axes=(0, 0, None))(positions, momentums, simulation_parameters)\n",
        "hamiltonians[:10]"
      ],
      "metadata": {
        "id": "fqGMEEUuPXyr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "shm_simulation.plot_coordinates(positions, momentums,\n",
        "                                simulation_parameters,\n",
        "                                title='TRUE TRAJECTORIES')"
      ],
      "metadata": {
        "id": "3ibwFZ8HPXyr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Coupled Harmonic Oscillators\n",
        "\n",
        "Consider a system of $n$ particles with masses $m_1, m_2, ..., m_n$.\n",
        "Each particle is connected to a wall with a spring of spring constant $k$,\n",
        "and every pair of particles is connected to each other with a spring of spring constant $\\kappa$.\n",
        "Then, the equations of motion take the following form:\n",
        "$$\n",
        "m_1\\frac{\\partial^2 q_i}{\\partial t^2} = -(k + \\kappa) q_i + \\sum_{j \\neq i}\\kappa q_j\n",
        "$$\n",
        "for each $i \\in \\{1, 2 ..., n\\}$.\n",
        "\n",
        "To find the normal modes where all particles oscillate with the same angular frequency $\\omega$, we make the exponential ansatz:\n",
        "$$\n",
        "q_i = c_ie^{i\\omega t}\n",
        "$$\n",
        "for each $i \\in \\{1, 2 ..., n\\}$.\n",
        "\n",
        "This gives us:\n",
        "$$\n",
        "-\\omega^2c_i = -\\frac{k + \\kappa}{m_i}c_i + \\sum_{j \\neq i}\\frac{\\kappa}{m_i}c_j\n",
        "$$\n",
        "for each $i \\in \\{1, 2 ..., n\\}$.\n",
        "In a matrix form, the coefficients $c$ satisfy\n",
        "$$\n",
        "(M + \\omega^2I_n)c = 0\n",
        "$$\n",
        "where $M$ is the matrix such that\n",
        "$M_{ii} = -\\frac{k + \\kappa}{m_i}$ and $M_{ij} = \\frac{\\kappa}{m_i}$ for $i \\neq j$.\n",
        "\n",
        "Thus, the angular frequencies $\\omega$ are the square roots of the negative\n",
        "eigenvalues of $M$."
      ],
      "metadata": {
        "id": "hZyga3HkQDq6"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "simulation_parameters = {\n",
        "    'm': jnp.asarray([2., 10.]),\n",
        "    'k_wall': jnp.asarray([10., 4.]),\n",
        "    'k_pair': jnp.asarray(1.),\n",
        "    'A': jnp.asarray([3., 1.]),\n",
        "    'phi': jnp.asarray([jnp.pi/2, 0.])\n",
        "}"
      ],
      "metadata": {
        "id": "6uV3LisYezPX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Generate coordinates at multiple instants of time!\n",
        "times = jnp.arange(100) * 1\n",
        "positions, momentums = jax.vmap(chm_simulation.generate_canonical_coordinates, in_axes=(0, None), out_axes=0)(\n",
        "    times, simulation_parameters)"
      ],
      "metadata": {
        "id": "JSQPSUAbYlCT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Compute Hamiltonians at the different instants.\n",
        "hamiltonians = jax.vmap(chm_simulation.compute_hamiltonian, in_axes=(0, 0, None))(\n",
        "    positions, momentums, simulation_parameters)\n",
        "hamiltonians"
      ],
      "metadata": {
        "id": "-huSXDDCvscY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Our Approach\n",
        "\n",
        "We propose learning the transformation to and from action-angle coordinates\n",
        "as neural networks.\n",
        "\n",
        "$$\n",
        "(q_0, p_0) \\xrightarrow{Enc} (\\theta_0, I) \\xrightarrow{\\theta_t \\ = \\ \\theta_0 + t \\cdot \\dot{\\theta}} (\\theta_t, I) \\xrightarrow{Dec} (q_t, p_t)\n",
        "$$\n",
        "\n",
        "To be precise, our ActionAngleNetwork performs the following operations:\n",
        "* Given canonical coordinates $(q_0, p_0)$, encode them with $Enc$ to action-angle coordinates $(\\theta_0, I)$.\n",
        "* Compute the Hamiltonian in these coordinates $H = Ham(I)$ with a HamiltonianNet $Ham$. Take the gradient of $H$ with respect to $I$ to get the angular velocity $\\dot{\\theta}$:\n",
        "\n",
        "$$\n",
        "\\dot{\\theta} = \\frac{d\\theta}{dt} = \\frac{dH}{dI}.\n",
        "$$\n",
        "\n",
        "* In the current implementation, we directly parametrize $\\dot{\\theta} = f_\\phi(I)$.\n",
        "* To get the angles at time $t$, we simply need to multiply with the angular velocity:\n",
        "\n",
        "$$\n",
        "\\theta_t = \\theta_0 + t \\cdot \\dot{\\theta}\n",
        "$$\n",
        "\n",
        "* Convert action-angle coordinates $(\\theta_t, I)$ back to canonical coordinates $(q_t, p_t)$, via the decoder $Dec$.\n",
        "\n",
        "We explore different parametrizations for the encoder $Enc$, HamiltonianNet $Ham$ and decoder \n",
        "$Dec$ neural networks.\n",
        "\n",
        "We also explore representing the action-angle space in Cartesian coordinates.\n",
        "In this formulation, the Encoder is followed by a fixed transformation to Polar coordinates.\n",
        "$$\n",
        "(q_0, p_0)\n",
        "\\xrightarrow{Enc} (\\hat{q}_0, \\hat{p}_0)\n",
        "\\xrightarrow{\\text{to Polar}} (\\theta_0, I)\n",
        "\\xrightarrow{\\theta_t \\ = \\ \\theta_0 + t \\cdot \\dot{\\theta}} (\\theta_t, I) \n",
        "\\xrightarrow{\\text{to Cartesian}} (\\hat{q}_t, \\hat{p}_t)\n",
        "\\xrightarrow{Dec} (q_t, p_t)\n",
        "$$"
      ],
      "metadata": {
        "id": "px2H9Bk-Aq_n"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Training Configuration\n",
        "config = action_angle_flow.get_config()\n",
        "workdir = tempfile.mkdtemp()"
      ],
      "metadata": {
        "id": "djEwafAOPjNH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%tensorboard --logdir={workdir} --port=0"
      ],
      "metadata": {
        "id": "E-c3sScK4A01"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "scaler, state, aux = train.train_and_evaluate(config, workdir)"
      ],
      "metadata": {
        "id": "2yAP_WC4DUYK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_positions = aux['train']['positions']\n",
        "train_momentums = aux['train']['momentums']\n",
        "train_simulation_parameters = aux['train']['simulation_parameters']\n",
        "all_train_metrics = aux['train']['metrics']\n",
        "\n",
        "test_positions = aux['test']['positions']\n",
        "test_momentums = aux['test']['momentums']\n",
        "test_simulation_parameters = aux['test']['simulation_parameters']\n",
        "all_test_metrics = aux['test']['metrics']"
      ],
      "metadata": {
        "id": "fSSkCqDmqMe1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Plotting Loss\n",
        "\n"
      ],
      "metadata": {
        "id": "XB3IRcoybC-G"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%config InlineBackend.figure_format = 'retina'\n",
        "\n",
        "metrics = all_test_metrics\n",
        "total_losses = {jump: np.asarray(list(metrics[step][jump]['prediction_loss'] for step in metrics)) for jump in config.test_time_jumps}\n",
        "steps = list(metrics.keys())\n",
        "colors = plt.cm.viridis(np.linspace(0, 1, len(total_losses)))\n",
        "\n",
        "fig, ax = plt.subplots()\n",
        "for jump_color, jump in zip(colors, config.test_time_jumps):\n",
        "  total_losses_for_jump = total_losses[jump]\n",
        "  ax.plot(steps, total_losses_for_jump, label=jump, color=jump_color)\n",
        "\n",
        "ax.set_title('Time Jump Sizes: Test Loss')\n",
        "ax.set_xlabel('Steps')\n",
        "ax.set_ylabel('Loss')\n",
        "ax.set_yscale('log')\n",
        "# ax.set_xscale('log')\n",
        "ax.legend(title='Jump Size')\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "Mj4iuzyxbCii"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "fig.savefig('test_losses.pdf', dpi=1000)\n",
        "files.download('test_losses.pdf')"
      ],
      "metadata": {
        "id": "HTQsINmOtNXh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Plotting Change in Hamiltonians"
      ],
      "metadata": {
        "id": "HxKQ4dYHJ8mM"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%config InlineBackend.figure_format = 'retina'\n",
        "\n",
        "metrics = all_test_metrics\n",
        "total_changes = {jump: np.asarray(list(metrics[step][jump]['mean_change_in_hamiltonians'] for step in metrics)) for jump in config.test_time_jumps}\n",
        "steps = list(metrics.keys())\n",
        "colors = plt.cm.viridis(np.linspace(0, 1, len(total_changes)))\n",
        "\n",
        "fig, ax = plt.subplots()\n",
        "for jump_color, jump in zip(colors, config.test_time_jumps):\n",
        "  total_changes_for_jump = total_changes[jump]\n",
        "  ax.plot(steps, total_changes_for_jump, label=jump, color=jump_color)\n",
        "\n",
        "true_position, true_momentum = train.inverse_transform_with_scaler(test_positions[0, :1], test_momentums[0, :1], scaler)\n",
        "actual_hamiltonian = shm_simulation.compute_hamiltonian(true_position, true_momentum, simulation_parameters)\n",
        "actual_hamiltonian = np.asarray(actual_hamiltonian).squeeze()\n",
        "ax.axhline(y=actual_hamiltonian, c='gray', linestyle='--')\n",
        "ax.set_title('Time Jump Sizes: Mean Change in Hamiltonian')\n",
        "ax.set_xlabel('Steps')\n",
        "ax.set_ylabel('Change')\n",
        "ax.set_yscale('log')\n",
        "# ax.set_xscale('log')\n",
        "ax.legend(title='Jump Size')\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "G3kioM9tJ8mM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "fig.savefig('test_change_in_hamiltonian.pdf', dpi=1000)\n",
        "files.download('test_change_in_hamiltonian.pdf')"
      ],
      "metadata": {
        "id": "ihgw9aguJ8mM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Plotting Action and Angle Space"
      ],
      "metadata": {
        "id": "_5NymBm6s6EG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Sample position-momentum space.\n",
        "max_position = 1.2 * np.abs(train_positions).max()\n",
        "max_momentum = 1.2 * np.abs(train_momentums).max()\n",
        "plot_positions = jnp.linspace(-max_position, max_position, num=100)\n",
        "plot_momentums = jnp.linspace(-max_momentum, max_momentum, num=100)\n",
        "grid = jnp.meshgrid(plot_positions, plot_momentums)\n",
        "plot_positions = grid[0][:, :, jnp.newaxis]\n",
        "plot_momentums = grid[1][:, :, jnp.newaxis]\n",
        "\n",
        "# Pad the remaining coordinates with zeros.\n",
        "def pad_coords(positions: chex.Array, momentums: chex.Array, index: int) -> Tuple[chex.Array, chex.Array]:\n",
        "  positions = jnp.pad(positions, ((0, 0), (0, 0), (index, config.num_trajectories - index - 1)))\n",
        "  momentums = jnp.pad(momentums, ((0, 0), (0, 0), (index, config.num_trajectories - index - 1)))\n",
        "  return positions, momentums\n",
        "\n",
        "trajectory_index = 1\n",
        "plot_positions, plot_momentums = pad_coords(plot_positions, plot_momentums, trajectory_index)\n",
        "\n",
        "# Compute actions and angles.\n",
        "_, _, auxiliary_predictions = jax.vmap(state.apply_fn, in_axes=(None, 0, 0, None))(state.params, plot_positions, plot_momentums, 0)\n",
        "plot_actions = auxiliary_predictions['actions']\n",
        "plot_angles = auxiliary_predictions['current_angles']\n",
        "\n",
        "train_positions_rescaled, train_momentums_rescaled = train.inverse_transform_with_scaler(train_positions, train_momentums, scaler)\n",
        "plot_positions, plot_momentums, plot_actions, plot_angles = jax.tree_map(lambda arr: arr[:, :, trajectory_index], (plot_positions, plot_momentums, plot_actions, plot_angles))\n",
        "plot_positions, plot_momentums = train.inverse_transform_with_scaler(plot_positions, plot_momentums, scaler)"
      ],
      "metadata": {
        "id": "nHq7EuKotBXG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "fig, ax = plt.subplots()\n",
        "contours = ax.contour(plot_positions, plot_momentums, plot_actions, 50, cmap='viridis')\n",
        "fig.colorbar(contours)\n",
        "ax.plot(train_positions_rescaled[:, trajectory_index], train_momentums_rescaled[:, trajectory_index],\n",
        "        c='gray', linestyle='--')\n",
        "ax.set_xlabel('q')\n",
        "ax.set_ylabel('p')\n",
        "ax.set_title('Actions Contour')\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "GbvA6Ljzxwys"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "fig.savefig('actions_contour.pdf', dpi=1000)\n",
        "files.download('actions_contour.pdf')"
      ],
      "metadata": {
        "id": "s41J4uCMtLDh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "fig, ax = plt.subplots()\n",
        "contours = ax.contour(plot_positions, plot_momentums, plot_angles, 50, cmap='viridis')\n",
        "fig.colorbar(contours)\n",
        "ax.plot(train_positions_rescaled[:, trajectory_index], train_momentums_rescaled[:, trajectory_index],\n",
        "        c='gray', linestyle='--')\n",
        "ax.set_xlabel('q')\n",
        "ax.set_ylabel('p')\n",
        "ax.set_title('Angles Contour')\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "LYdLzUSMyqjf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "fig.savefig('angles_contour.pdf', dpi=1000)\n",
        "files.download('angles_contour.pdf')"
      ],
      "metadata": {
        "id": "in1QA6s_tOnc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Plotting True Trajectories"
      ],
      "metadata": {
        "id": "RQDttGU2zrRU"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "train_positions_rescaled, train_momentums_rescaled = train.inverse_transform_with_scaler(train_positions, train_momentums, scaler)\n",
        "max_position = np.abs(train_positions_rescaled).max()\n",
        "max_momentum = np.abs(train_momentums_rescaled).max()\n",
        "shm_simulation.static_plot_coordinates_in_phase_space(train_positions_rescaled, train_momentums_rescaled, title='TRAIN TRAJECTORIES', max_position=max_position, max_momentum=max_momentum)"
      ],
      "metadata": {
        "id": "Jdvko_Omzq55"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "test_positions_rescaled, test_momentums_rescaled = train.inverse_transform_with_scaler(test_positions, test_momentums, scaler)\n",
        "max_position = np.abs(train_positions_rescaled).max()\n",
        "max_momentum = np.abs(train_momentums_rescaled).max()\n",
        "shm_simulation.static_plot_coordinates_in_phase_space(test_positions_rescaled, test_momentums_rescaled, title='TEST TRAJECTORIES', max_position=max_position, max_momentum=max_momentum)"
      ],
      "metadata": {
        "id": "8DN4DZ-5PvjU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### One-step Predictions"
      ],
      "metadata": {
        "id": "k-_VH5inpan4"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def predict_for_trajectory(positions_for_trajectory: chex.Array, momentums_for_trajectory: chex.Array, jump: int):\n",
        "  curr_positions, curr_momentums, target_positions, target_momentums = train.get_coordinates_for_time_jump(positions_for_trajectory, momentums_for_trajectory, jump)\n",
        "  predicted_positions, predicted_momentums, auxiliary_predictions = train.compute_predictions(state, curr_positions, curr_momentums, jump * config.time_delta)\n",
        "  predicted_positions, predicted_momentums = train.inverse_transform_with_scaler(predicted_positions, predicted_momentums, scaler)\n",
        "  return predicted_positions, predicted_momentums"
      ],
      "metadata": {
        "id": "ERuwQNd0htS-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "jump = 1\n",
        "train_positions_rescaled, train_momentums_rescaled = train.inverse_transform_with_scaler(train_positions, train_momentums, scaler)\n",
        "max_position = np.abs(train_positions_rescaled).max()\n",
        "max_momentum = np.abs(train_momentums_rescaled).max()\n",
        "\n",
        "predicted_positions, predicted_momentums = predict_for_trajectory(train_positions, train_momentums, jump)\n",
        "shm_simulation.static_plot_coordinates_in_phase_space(predicted_positions, predicted_momentums, title=f'PREDICTED TRAIN TRAJECTORIES: JUMP {jump}', max_position=max_position, max_momentum=max_momentum)"
      ],
      "metadata": {
        "id": "TqMXFbrYdJ74"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_one_step_predicted_trajectories_anim = shm_simulation.plot_coordinates(\n",
        "    predicted_positions[:200, 0], predicted_momentums[:200, 0], jax.tree_map(lambda arr: arr[0], train_simulation_parameters),\n",
        "    title=f'ONE-STEP PREDICTED TRAIN TRAJECTORIES: JUMP {jump}')\n",
        "train_one_step_predicted_trajectories_anim"
      ],
      "metadata": {
        "id": "_RfpypoXdMr4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "jump = 200\n",
        "train_positions_rescaled, train_momentums_rescaled = train.inverse_transform_with_scaler(train_positions, train_momentums, scaler)\n",
        "max_position = np.abs(train_positions_rescaled).max()\n",
        "max_momentum = np.abs(train_momentums_rescaled).max()\n",
        "\n",
        "predicted_positions, predicted_momentums = predict_for_trajectory(test_positions, test_momentums, jump)\n",
        "shm_simulation.static_plot_coordinates_in_phase_space(predicted_positions, predicted_momentums,\n",
        "                                                      title=f'PREDICTED TEST TRAJECTORIES: JUMP {jump}',\n",
        "                                                      max_position=max_position, max_momentum=max_momentum)"
      ],
      "metadata": {
        "id": "-XsXCNhiuG2q"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "test_one_step_predicted_trajectories_anim = shm_simulation.plot_coordinates(\n",
        "    predicted_positions[:200, 0], predicted_momentums[:200, 0], jax.tree_map(lambda arr: arr[0], test_simulation_parameters),\n",
        "    title=f'ONE-STEP PREDICTED TEST TRAJECTORIES: JUMP {jump}')\n",
        "test_one_step_predicted_trajectories_anim"
      ],
      "metadata": {
        "id": "_hmg_H0cq9jO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "test_one_step_predicted_trajectories_anim.save('test_one_step_predicted_trajectories.gif')\n",
        "files.download('test_one_step_predicted_trajectories.gif')"
      ],
      "metadata": {
        "id": "P_gTuJAtKbeB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "test_one_step_predicted_trajectories_phase_space_anim = shm_simulation.plot_coordinates_in_phase_space(\n",
        "    predicted_positions[:100], predicted_momentums[:100],\n",
        "    test_simulation_parameters,\n",
        "    title=f'ONE-STEP PREDICTED TEST TRAJECTORIES: JUMP {jump}')\n",
        "test_one_step_predicted_trajectories_phase_space_anim"
      ],
      "metadata": {
        "id": "9bRDoATiS3DV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "test_one_step_predicted_trajectories_phase_space_anim.save('test_one_step_predicted_trajectories_phase_space.gif')\n",
        "files.download('test_one_step_predicted_trajectories_phase_space.gif')"
      ],
      "metadata": {
        "id": "SQxrPejKLlJI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Distribution of Actions"
      ],
      "metadata": {
        "id": "YLHw4DZwmA4T"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Compute actions.\n",
        "jump = 1\n",
        "curr_positions, curr_momentums, *_ = train.get_coordinates_for_time_jump(train_positions[0], train_momentums[0], jump)\n",
        "\n",
        "_, _, auxiliary_predictions = state.apply_fn(state.params, curr_positions, curr_momentums, 0)\n",
        "plot_actions = auxiliary_predictions['actions']\n",
        "plot_angles = auxiliary_predictions['curr_angles']\n",
        "\n",
        "actions = auxiliary_predictions['actions']\n",
        "actions = np.asarray(actions).flatten()\n",
        "plt.hist(actions, bins=1000)\n",
        "plt.gca().xaxis.set_major_formatter(matplotlib.ticker.StrMethodFormatter('{x:,.3f}'))\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "SqASSBMew2uP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "angular_velocities = auxiliary_predictions['angular_velocities']\n",
        "angular_velocities = np.asarray(angular_velocities).flatten()\n",
        "plt.hist(angular_velocities, bins=1000)\n",
        "plt.gca().xaxis.set_major_formatter(matplotlib.ticker.StrMethodFormatter('{x:,.4f}'))\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "NYutCNyMfgLe"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Non-Recursive Multi-Step Predictions"
      ],
      "metadata": {
        "id": "AnOJUuacpf0C"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "jump = 50\n",
        "non_recursive_predictions = state.apply_fn(\n",
        "    state.params, test_positions[:1], test_momentums[:1],\n",
        "    jnp.arange(1, 500) * jump * config.time_delta, method=models.ActionAngleNetwork.predict_multi_step)\n",
        "non_recursive_multi_step_predicted_positions, non_recursive_multi_step_predicted_momentums, non_recursive_multi_step_auxiliary_predictions = non_recursive_predictions\n",
        "non_recursive_multi_step_auxiliary_predictions['angular_velocities'], non_recursive_multi_step_auxiliary_predictions['current_angles'], non_recursive_multi_step_auxiliary_predictions['future_angles'][:10]"
      ],
      "metadata": {
        "id": "K0c2cmbV4K6-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "non_recursive_multi_step_predicted_positions, non_recursive_multi_step_predicted_momentums = train.inverse_transform_with_scaler(non_recursive_multi_step_predicted_positions, non_recursive_multi_step_predicted_momentums, scaler)"
      ],
      "metadata": {
        "id": "FUAEgpRfje5V"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "fig = shm_simulation.static_plot_coordinates_in_phase_space(\n",
        "    non_recursive_multi_step_predicted_positions[:100], non_recursive_multi_step_predicted_momentums[:100],\n",
        "    title=f'PREDICTED TEST TRAJECTORIES: JUMP {jump}')\n",
        "fig"
      ],
      "metadata": {
        "id": "ZsFM_JZDQAeE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "fig.savefig(f'test_multi_step_predicted_phase_space_jump_{jump}.pdf', dpi=1000, bbox_inches='tight', pad_inches=0)\n",
        "files.download('test_multi_step_predicted_phase_space_jump_{jump}.pdf')"
      ],
      "metadata": {
        "id": "xe_HVs17uTzh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "test_multi_step_predicted_trajectories_anim = plot_coordinates(\n",
        "    non_recursive_multi_step_predicted_positions, non_recursive_multi_step_predicted_momentums, simulation_parameters,\n",
        "    title='NON-RECURSIVE MULTI-STEP PREDICTED TEST TRAJECTORIES')\n",
        "test_multi_step_predicted_trajectories_anim"
      ],
      "metadata": {
        "id": "1U9xwtDWRDo8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "test_multi_step_predicted_trajectories_anim.save('test_non_recursive_multi_step_predicted_trajectories.gif')\n",
        "files.download('test_non_recursive_multi_step_predicted_trajectories.gif')"
      ],
      "metadata": {
        "id": "wukONGNXRDpD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "test_multi_step_predicted_trajectories_phase_space_anim = shm_simulation.plot_coordinates_in_phase_space(\n",
        "    non_recursive_multi_step_predicted_positions[:100], non_recursive_multi_step_predicted_momentums[:100], test_simulation_parameters,\n",
        "    title=f'PREDICTED TEST TRAJECTORIES: JUMP {jump}')\n",
        "test_multi_step_predicted_trajectories_phase_space_anim"
      ],
      "metadata": {
        "id": "nbdBYsHfRDpD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "test_multi_step_predicted_trajectories_phase_space_anim.save(f'test_non_recursive_multi_step_predicted_trajectories_phase_space_jump_{jump}.gif')\n",
        "files.download(f'test_non_recursive_multi_step_predicted_trajectories_phase_space_jump_{jump}.gif')"
      ],
      "metadata": {
        "id": "R5VVTn4BRDpD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Recursive Multi-Step Predictions"
      ],
      "metadata": {
        "id": "H09MhjHiReRL"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def predict_next_step(carry, _):\n",
        "  current_position, current_momentum = carry\n",
        "  predicted_position, predicted_momentum, auxiliary_predictions = train.compute_predictions(\n",
        "      state, current_position, current_momentum, config.time_delta)\n",
        "  return (predicted_position, predicted_momentum), (predicted_position, predicted_momentum, auxiliary_predictions)\n",
        "\n",
        "_, recursive_predictions = jax.lax.scan(predict_next_step, (test_positions[0, :1], test_momentums[0, :1]), None, length=1000)\n",
        "recursive_multi_step_predicted_positions, recursive_multi_step_predicted_momentums, recursive_multi_step_auxiliary_predictions = recursive_predictions\n",
        "recursive_multi_step_auxiliary_predictions['angular_velocities'][0], recursive_multi_step_auxiliary_predictions['current_angles'][:10], recursive_multi_step_auxiliary_predictions['future_angles'][:10]"
      ],
      "metadata": {
        "id": "3srAj0PxXUR4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "recursive_multi_step_predicted_positions[0], recursive_multi_step_predicted_momentums[0]"
      ],
      "metadata": {
        "id": "Kg7wdSjgHW1C"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "recursive_multi_step_predicted_positions, recursive_multi_step_predicted_momentums = train.inverse_transform_with_scaler(\n",
        "    recursive_multi_step_predicted_positions, recursive_multi_step_predicted_momentums, scaler)"
      ],
      "metadata": {
        "id": "fmUazZnWlvNZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "recursive_multi_step_predicted_positions[0], recursive_multi_step_predicted_momentums[0]"
      ],
      "metadata": {
        "id": "hHox2qP4PSwO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "test_multi_step_predicted_trajectories_anim = shm_simulation.plot_coordinates(\n",
        "    recursive_multi_step_predicted_positions[:200], recursive_multi_step_predicted_momentums[:200], simulation_parameters,\n",
        "    title='RECURSIVE MULTI-STEP PREDICTED TEST TRAJECTORIES')\n",
        "test_multi_step_predicted_trajectories_anim"
      ],
      "metadata": {
        "id": "CnKVFzUIdCUD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "test_multi_step_predicted_trajectories_anim.save('test_recursive_multi_step_predicted_trajectories.gif')\n",
        "files.download('test_recursive_multi_step_predicted_trajectories.gif')"
      ],
      "metadata": {
        "id": "Zd5qVxLOeNlr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "shm_simulation.static_plot_coordinates_in_phase_space(\n",
        "    recursive_multi_step_predicted_positions, recursive_multi_step_predicted_momentums,\n",
        "    title='RECURSIVE MULTI-STEP PREDICTED TEST TRAJECTORIES')"
      ],
      "metadata": {
        "id": "IvNO-958Ah9a"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "test_multi_step_predicted_trajectories_phase_space_anim = shm_simulation.plot_coordinates_in_phase_space(\n",
        "    recursive_multi_step_predicted_positions[:200], recursive_multi_step_predicted_momentums[:200], simulation_parameters,\n",
        "    title='RECURSIVE MULTI-STEP PREDICTED TEST TRAJECTORIES')\n",
        "test_multi_step_predicted_trajectories_phase_space_anim"
      ],
      "metadata": {
        "id": "6Fp2JjnBTntG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "test_multi_step_predicted_trajectories_phase_space_anim.save('test_recursive_multi_step_predicted_trajectories_phase_space.gif')\n",
        "files.download('test_recursive_multi_step_predicted_trajectories_phase_space.gif')"
      ],
      "metadata": {
        "id": "mItVG-inTntO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### True Trajectories"
      ],
      "metadata": {
        "id": "Tr73D2aZpdBT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "test_true_trajectories_anim = shm_simulation.plot_coordinates(\n",
        "    test_target_positions[:200], test_target_momentums[:200], simulation_parameters,\n",
        "    title='TRUE TEST TRAJECTORIES')\n",
        "test_true_trajectories_anim"
      ],
      "metadata": {
        "id": "4qARWdeFErrq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "test_true_trajectories_anim.save('test_true_trajectories.gif')\n",
        "files.download('test_true_trajectories.gif')"
      ],
      "metadata": {
        "id": "vFX_A7pkFldi"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}